#include <sys/mman.h>
#include <pthread.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>

static pthread_t alloc_mine;
static pthread_mutex_t * alloc_lock;
static pthread_cond_t * alloc_clear;
static unsigned int alloc_inuse;

static void enter_alloc(void)
{
	if (pthread_equal(pthread_self(), alloc_mine))
		return;
	
	pthread_mutex_lock(alloc_lock);
	alloc_inuse++;
	pthread_mutex_unlock(alloc_lock);
}

static void leave_alloc(void)
{
	if (pthread_equal(pthread_self(), alloc_mine))
		return;
	
	pthread_mutex_lock(alloc_lock);
	alloc_inuse--;
	if (alloc_clear != NULL &&
	    alloc_inuse == 0)
		pthread_cond_signal(alloc_clear);
	pthread_mutex_unlock(alloc_lock);
}

static void freeze_alloc(void)
{
	pthread_cond_t * clear;
	
	clear = malloc(sizeof(*clear));
	pthread_cond_init(clear, NULL);
	
	pthread_mutex_lock(alloc_lock);
	alloc_clear = clear;
	while (alloc_inuse != 0)
		pthread_cond_wait(alloc_clear, alloc_lock);
	alloc_clear = NULL;
	alloc_mine = pthread_self();
	
	pthread_cond_destroy(clear);
	free(clear);
}

static void resume_alloc(void)
{
	pthread_mutex_unlock(alloc_lock);
	alloc_mine = 0;
}

static void reinit_alloc(void)
{
	pthread_mutex_init(alloc_lock, NULL);
	alloc_mine = 0;
}

static int expected_program(void)
{
	FILE * file;
	char * line;
	size_t size;
	ssize_t read;
	int match;
	
	file = fopen("/proc/self/comm", "r");
	line = NULL;
	size = 0;
	
	read = getline(&line, &size, file);
	line[read - 1] = '\0';
	match = strcmp(line, "chrome");
	
	fclose(file);
	free(line);
	
	return match;
}

static const char headcode_type1[] = { 0x41, 0x57,
				       0x41, 0x56,
				       0x41, 0x55,
				       0x41, 0x54,
				       0x55,
				       0x53,
				       0x48, 0x89, 0xfb };
static const char headcode_type2[] = { 0x41, 0x57,
				       0x41, 0x56,
				       0x41, 0x55,
				       0x41, 0x54,
				       0x55,
				       0x53,
				       0x48, 0x89, 0xf3 };
static const char headcode_type3[] = { 0x48, 0x89, 0x5c, 0x24, 0xd0,
				       0x48, 0x89, 0x6c, 0x24, 0xd8,
				       0x48, 0x89, 0xfb };

static const char hookcode_part1[] = { 0x48, 0x89, 0x5c, 0x24, 0xf0,
				       0x48, 0x89, 0x6c, 0x24, 0xf8,
				       0x48, 0x89, 0xfb,
				       0x48, 0x89, 0xf5,
				       0x48, 0x83, 0xec, 0x18,
				       0x48, 0xb8 };
static const char hookcode_part2[] = { 0xff, 0xd0,
				       0x48, 0x89, 0xee,
				       0x48, 0x89, 0xdf,
				       0x48, 0xb8 };
static const char hookcode_part3[] = { 0xff, 0xd0,
				       0x48, 0x89, 0xc3,
				       0x48, 0xb8 };
static const char hookcode_part4[] = { 0xff, 0xd0,
				       0x48, 0x89, 0xd8,
				       0x48, 0x8b, 0x6c, 0x24, 0x10,
				       0x48, 0x8b, 0x5c, 0x24, 0x08,
				       0x48, 0x83, 0xc4, 0x18,
				       0xc3 };
static const char hookcode_part5[] = { 0x48, 0xb8 };
static const char hookcode_part6[] = { 0xff, 0xe0 };

static void override_function(void * baseaddr,
			      const void * headcode,
			      size_t headsize)
{
	size_t basesize;
	void * basepage;
	void * wrapaddr;
	void * instruct;
	
	if (memcmp(baseaddr, headcode, headsize) != 0) {
		fprintf(stderr, "%p is not what we expected\n", baseaddr);
		_exit(-1);
	}
	basesize = ((size_t)baseaddr & 0x0fff) + headsize;
	basepage = (void *)((size_t)baseaddr & ~0x0fff);
	
	wrapaddr = mmap(NULL, 0x1000, PROT_READ | PROT_WRITE | PROT_EXEC,
			MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
	mprotect(basepage, basesize, PROT_READ | PROT_WRITE | PROT_EXEC);
	
	instruct = wrapaddr;
	memcpy(instruct, hookcode_part1, sizeof(hookcode_part1));
	instruct += sizeof(hookcode_part1);
	*(void **)instruct = enter_alloc;
	instruct += sizeof(void *);
	memcpy(instruct, hookcode_part2, sizeof(hookcode_part2));
	instruct += sizeof(hookcode_part2);
	*(void **)instruct = wrapaddr + 0x60;
	instruct += sizeof(void *);
	memcpy(instruct, hookcode_part3, sizeof(hookcode_part3));
	instruct += sizeof(hookcode_part3);
	*(void **)instruct = leave_alloc;
	instruct += sizeof(void *);
	memcpy(instruct, hookcode_part4, sizeof(hookcode_part4));
	
	instruct = wrapaddr + 0x60;
	memcpy(instruct, headcode, headsize);
	instruct += headsize;
	memcpy(instruct, hookcode_part5, sizeof(hookcode_part5));
	instruct += sizeof(hookcode_part5);
	*(void **)instruct = baseaddr + headsize;
	instruct += sizeof(void *);
	memcpy(instruct, hookcode_part6, sizeof(hookcode_part6));
	
	instruct = baseaddr;
	memcpy(instruct, hookcode_part5, sizeof(hookcode_part5));
	instruct += sizeof(hookcode_part5);
	*(void **)instruct = wrapaddr;
	instruct += sizeof(void *);
	memcpy(instruct, hookcode_part6, sizeof(hookcode_part6));
	
	mprotect(basepage, basesize, PROT_READ | PROT_EXEC);
	mprotect(wrapaddr, 0x1000, PROT_READ | PROT_EXEC);
}

static void __attribute__((constructor)) override_allocator(void)
{
	pthread_mutex_t * lock;
	
	if (expected_program())
		return;
	
	pthread_atfork(freeze_alloc,
		       resume_alloc,
		       reinit_alloc);
	lock = malloc(sizeof(*lock));
	pthread_mutex_init(lock, NULL);
	alloc_mine = 0;
	alloc_lock = lock;
	alloc_clear = NULL;
	alloc_inuse = 0;
	
	override_function(malloc,
			  headcode_type1,
			  sizeof(headcode_type1));
	override_function(calloc,
			  headcode_type2,
			  sizeof(headcode_type2));
	override_function(realloc,
			  headcode_type3,
			  sizeof(headcode_type3));
	override_function(free,
			  headcode_type3,
			  sizeof(headcode_type3));
}
